#!/usr/bin/env python
# -*- coding: utf-8 -*-

from PyQt5.QtWidgets import (QWidget, QVBoxLayout, QHBoxLayout, QPushButton, 
                            QListWidget, QListWidgetItem, QLabel, QComboBox,
                            QToolBar, QAction, QFileDialog, QSplitter, 
                            QScrollArea, QFrame, QGroupBox, QFormLayout, 
                            QLineEdit, QSpinBox, QDoubleSpinBox, QCheckBox,
                            QTabWidget, QTextEdit, QSlider, QRadioButton,
                            QButtonGroup, QMessageBox, QProgressBar, QGridLayout,
                            QSizePolicy, QFileDialog)
from PyQt5.QtCore import Qt, QSize, QRect, QPoint, pyqtSignal, QThread, QObject, QTimer, QUrl
from PyQt5.QtGui import QIcon, QPixmap, QPainter, QPen, QColor, QImage, QFont
import os
import cv2
import glob
import numpy as np
from pathlib import Path
import time

# 检查是否能导入ultralytics库
try:
    from ultralytics import YOLO
    ULTRALYTICS_AVAILABLE = True
except ImportError:
    ULTRALYTICS_AVAILABLE = False

class InferenceThread(QThread):
    """推理线程类"""
    progress_updated = pyqtSignal(int)  # 进度更新信号
    result_ready = pyqtSignal(object, object)  # 结果和原始图像信号
    inference_finished = pyqtSignal(bool, str)  # 完成信号：成功状态，消息
    inference_stats = pyqtSignal(dict)  # 推理统计信号
    video_frame = pyqtSignal(object, object)  # 视频帧信号：结果，原始图像
    
    def __init__(self, model_path, source, conf_thres=0.25, iou_thres=0.45, classes=None, source_type="image"):
        super().__init__()
        self.model_path = model_path
        self.source = source
        self.conf_thres = conf_thres
        self.iou_thres = iou_thres
        self.classes = classes
        self.running = True
        self.source_type = source_type  # "image", "video", "camera"
        
    def run(self):
        """运行推理"""
        try:
            if not ULTRALYTICS_AVAILABLE:
                self.inference_finished.emit(False, "未安装ultralytics库，请先安装")
                return
                
            # 检查模型文件是否存在
            if not os.path.exists(self.model_path):
                self.inference_finished.emit(False, f"模型文件不存在: {self.model_path}")
                return
            
            # 检查源文件或摄像头
            if self.source_type in ["image", "video"] and not os.path.exists(self.source):
                self.inference_finished.emit(False, f"源文件不存在: {self.source}")
                return
                
            # 加载模型
            self.progress_updated.emit(10)
            model = YOLO(self.model_path)
            self.progress_updated.emit(30)
            
            # 根据输入源类型处理
            if self.source_type == "image":
                # 图像推理
                self._process_image(model)
            elif self.source_type == "video":
                # 视频推理
                self._process_video(model)
            elif self.source_type == "camera":
                # 摄像头推理
                self._process_camera(model)
            else:
                self.inference_finished.emit(False, f"不支持的输入源类型: {self.source_type}")
                
        except Exception as e:
            import traceback
            error_msg = f"推理失败: {str(e)}\n{traceback.format_exc()}"
            self.inference_finished.emit(False, error_msg)
            
    def _process_image(self, model):
        """处理图像推理"""
        start_time = time.time()
        results = model(self.source, 
                       conf=self.conf_thres, 
                       iou=self.iou_thres,
                       classes=self.classes)
        inference_time = time.time() - start_time
        
        # 读取原始图像
        orig_img = cv2.imread(self.source)
        orig_img = cv2.cvtColor(orig_img, cv2.COLOR_BGR2RGB)
        
        # 发送结果
        self.result_ready.emit(results[0], orig_img)
        
        # 计算统计信息
        stats = {
            'inference_time': inference_time,
            'detected_objects': len(results[0].boxes),
            'image_size': f"{orig_img.shape[1]}x{orig_img.shape[0]}"
        }
        self.inference_stats.emit(stats)
        
        self.progress_updated.emit(100)
        self.inference_finished.emit(True, "推理完成")
    
    def _process_video(self, model):
        """处理视频推理"""
        try:
            # 打开视频文件
            cap = cv2.VideoCapture(self.source)
            if not cap.isOpened():
                self.inference_finished.emit(False, f"无法打开视频文件: {self.source}")
                return
                
            # 获取视频信息
            fps = cap.get(cv2.CAP_PROP_FPS)
            total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
            width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
            height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
            
            # 更新进度条
            self.progress_updated.emit(40)
            
            # 初始化统计信息
            total_inference_time = 0
            frame_count = 0
            total_objects = 0
            
            # 处理每一帧
            while self.running and cap.isOpened():
                ret, frame = cap.read()
                if not ret:
                    break
                    
                # 更新进度
                if total_frames > 0:
                    progress = int(40 + (frame_count / total_frames) * 60)
                    self.progress_updated.emit(progress)
                    
                # 处理帧
                start_time = time.time()
                results = model(frame, conf=self.conf_thres, iou=self.iou_thres, classes=self.classes)
                inference_time = time.time() - start_time
                
                # 转换颜色空间
                frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                
                # 更新统计信息
                total_inference_time += inference_time
                frame_count += 1
                total_objects += len(results[0].boxes)
                
                # 发送结果帧
                self.video_frame.emit(results[0], frame_rgb)
                
                # 限制帧率，避免UI过载
                time.sleep(max(0, 1/30 - inference_time))
                
            # 关闭视频
            cap.release()
            
            # 发送最终统计信息
            if frame_count > 0:
                avg_inference_time = total_inference_time / frame_count
                avg_objects = total_objects / frame_count
                
                stats = {
                    'inference_time': avg_inference_time,
                    'detected_objects': avg_objects,
                    'image_size': f"{width}x{height}",
                    'fps': fps,
                    'frames': frame_count
                }
                self.inference_stats.emit(stats)
                
            self.progress_updated.emit(100)
            self.inference_finished.emit(True, f"视频推理完成，共处理 {frame_count} 帧")
            
        except Exception as e:
            import traceback
            error_msg = f"视频推理失败: {str(e)}\n{traceback.format_exc()}"
            self.inference_finished.emit(False, error_msg)
            
    def _process_camera(self, model):
        """处理摄像头推理"""
        try:
            # 打开摄像头
            camera_id = 0
            if self.source.isdigit():
                camera_id = int(self.source)
                
            cap = cv2.VideoCapture(camera_id)
            if not cap.isOpened():
                self.inference_finished.emit(False, f"无法打开摄像头 ID: {camera_id}")
                return
                
            # 获取摄像头信息
            width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
            height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
            
            # 更新进度条
            self.progress_updated.emit(50)
            
            # 初始化统计信息
            total_inference_time = 0
            frame_count = 0
            total_objects = 0
            
            # 处理每一帧
            while self.running and cap.isOpened():
                ret, frame = cap.read()
                if not ret:
                    break
                    
                # 处理帧
                start_time = time.time()
                results = model(frame, conf=self.conf_thres, iou=self.iou_thres, classes=self.classes)
                inference_time = time.time() - start_time
                
                # 转换颜色空间
                frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                
                # 更新统计信息
                total_inference_time += inference_time
                frame_count += 1
                total_objects += len(results[0].boxes)
                
                # 计算并显示当前FPS
                current_fps = 1.0 / max(inference_time, 0.001)
                
                # 每10帧更新一次统计信息
                if frame_count % 10 == 0:
                    stats = {
                        'inference_time': inference_time,
                        'detected_objects': len(results[0].boxes),
                        'image_size': f"{width}x{height}",
                        'fps': current_fps
                    }
                    self.inference_stats.emit(stats)
                
                # 发送结果帧
                self.video_frame.emit(results[0], frame_rgb)
                
                # 限制帧率，避免UI过载
                time.sleep(max(0, 1/30 - inference_time))
                
            # 关闭摄像头
            cap.release()
            
            # 发送最终统计信息
            if frame_count > 0:
                avg_inference_time = total_inference_time / frame_count
                avg_objects = total_objects / frame_count
                
                stats = {
                    'inference_time': avg_inference_time,
                    'detected_objects': avg_objects,
                    'image_size': f"{width}x{height}",
                    'fps': frame_count / max(total_inference_time, 0.001),
                    'frames': frame_count
                }
                self.inference_stats.emit(stats)
                
            self.progress_updated.emit(100)
            self.inference_finished.emit(True, f"摄像头推理已停止，共处理 {frame_count} 帧")
            
        except Exception as e:
            import traceback
            error_msg = f"摄像头推理失败: {str(e)}\n{traceback.format_exc()}"
            self.inference_finished.emit(False, error_msg)
            
    def stop(self):
        """停止推理"""
        self.running = False

class ResultsWidget(QLabel):
    """结果显示窗口部件"""
    
    def __init__(self):
        super().__init__()
        self.setAlignment(Qt.AlignCenter)
        self.setMinimumSize(640, 480)
        self.setFrameStyle(QFrame.StyledPanel)
        self.setScaledContents(False)
        self.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
        self.results = None
        self.orig_image = None
        
    def set_results(self, results, orig_image):
        """设置结果和原始图像"""
        self.results = results
        self.orig_image = orig_image
        self.update_display()
        
    def update_display(self):
        """更新显示内容"""
        if self.results is None or self.orig_image is None:
            return
            
        # 创建带标注框的图像
        if hasattr(self.results, 'plot'):
            # 使用ultralytics的绘图功能
            result_img = self.results.plot()
            result_img = cv2.cvtColor(result_img, cv2.COLOR_BGR2RGB)
            
            # 将NumPy数组转换为QImage
            h, w, c = result_img.shape
            q_img = QImage(result_img.data, w, h, w * c, QImage.Format_RGB888)
            
            # 转换为QPixmap并显示
            pixmap = QPixmap.fromImage(q_img)
            
            # 适应窗口大小
            self.setPixmap(pixmap.scaled(self.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation))
        else:
            # 没有结果对象，显示原始图像
            h, w, c = self.orig_image.shape
            q_img = QImage(self.orig_image.data, w, h, w * c, QImage.Format_RGB888)
            pixmap = QPixmap.fromImage(q_img)
            self.setPixmap(pixmap.scaled(self.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation))
        
    def resizeEvent(self, event):
        """窗口大小变化时调整图像"""
        super().resizeEvent(event)
        if self.pixmap() is not None:
            self.update_display()

class InferencePanel(QWidget):
    """模型推理面板"""
    
    def __init__(self):
        super().__init__()
        self.inference_thread = None
        self.init_ui()
        
    def init_ui(self):
        """初始化UI"""
        main_layout = QVBoxLayout(self)
        main_layout.setContentsMargins(20, 20, 20, 20)
        main_layout.setSpacing(20)
        
        # 创建分割器
        splitter = QSplitter(Qt.Horizontal)
        
        # 左侧控制面板
        left_panel = QWidget()
        left_layout = QVBoxLayout(left_panel)
        left_layout.setContentsMargins(0, 0, 0, 0)
        left_layout.setSpacing(20)
        
        # 标题
        title_label = QLabel("模型推理")
        title_label.setProperty("title", "true")
        title_font = QFont()
        title_font.setPointSize(16)
        title_font.setBold(True)
        title_label.setFont(title_font)
        left_layout.addWidget(title_label)
        
        # 模型选择组
        model_group = QGroupBox("模型选择")
        model_layout = QVBoxLayout(model_group)
        
        # 模型路径选择
        model_path_layout = QHBoxLayout()
        self.model_path_edit = QLineEdit()
        self.model_path_edit.setReadOnly(True)
        self.model_path_edit.setPlaceholderText("选择YOLO模型文件")
        
        browse_model_btn = QPushButton("浏览...")
        browse_model_btn.setFixedWidth(100)
        browse_model_btn.clicked.connect(self.browse_model)
        
        model_path_layout.addWidget(self.model_path_edit)
        model_path_layout.addWidget(browse_model_btn)
        
        model_layout.addLayout(model_path_layout)
        
        # 模型下拉选择框
        model_combo_layout = QHBoxLayout()
        model_combo_layout.addWidget(QLabel("选择已有模型:"))
        
        self.model_combo = QComboBox()
        self.refresh_model_list()
        self.model_combo.currentIndexChanged.connect(self.on_model_selected)
        
        model_combo_layout.addWidget(self.model_combo)
        
        refresh_btn = QPushButton("刷新")
        refresh_btn.setFixedWidth(80)
        refresh_btn.clicked.connect(self.refresh_model_list)
        model_combo_layout.addWidget(refresh_btn)
        
        model_layout.addLayout(model_combo_layout)
        
        left_layout.addWidget(model_group)
        
        # 输入源选择组
        source_group = QGroupBox("输入源")
        source_layout = QVBoxLayout(source_group)
        
        # 源类型选择
        source_type_layout = QHBoxLayout()
        source_type_layout.addWidget(QLabel("源类型:"))
        self.source_type_combo = QComboBox()
        self.source_type_combo.addItems(["图像", "视频", "摄像头"])
        self.source_type_combo.currentIndexChanged.connect(self.on_source_type_changed)
        source_type_layout.addWidget(self.source_type_combo)
        source_layout.addLayout(source_type_layout)
        
        # 源文件选择
        source_path_layout = QHBoxLayout()
        self.source_path_edit = QLineEdit()
        self.source_path_edit.setReadOnly(True)
        self.source_path_edit.setPlaceholderText("选择图像或视频文件")
        
        self.browse_source_btn = QPushButton("浏览...")
        self.browse_source_btn.setFixedWidth(100)
        self.browse_source_btn.clicked.connect(self.browse_source)
        
        source_path_layout.addWidget(self.source_path_edit)
        source_path_layout.addWidget(self.browse_source_btn)
        
        source_layout.addLayout(source_path_layout)
        
        # 摄像头选择（默认隐藏）
        camera_layout = QHBoxLayout()
        camera_label = QLabel("摄像头ID:")
        camera_layout.addWidget(camera_label)
        self.camera_id_combo = QComboBox()
        self.camera_id_combo.addItems(["0", "1", "2", "3"])
        self.camera_id_combo.setCurrentIndex(0)
        camera_layout.addWidget(self.camera_id_combo)
        self.camera_layout = camera_layout  # 保存引用以便控制显示/隐藏
        source_layout.addLayout(camera_layout)
        
        # 默认隐藏摄像头选择
        self.camera_id_combo.setVisible(False)
        camera_label.setVisible(False)  # 使用直接引用
        self.camera_label = camera_label  # 保存引用以便后续使用
        
        left_layout.addWidget(source_group)
        
        # 推理参数组
        params_group = QGroupBox("推理参数")
        params_layout = QFormLayout(params_group)
        
        # 置信度阈值
        self.conf_threshold = QDoubleSpinBox()
        self.conf_threshold.setRange(0.01, 1.0)
        self.conf_threshold.setSingleStep(0.05)
        self.conf_threshold.setValue(0.25)
        params_layout.addRow("置信度阈值:", self.conf_threshold)
        
        # IoU阈值
        self.iou_threshold = QDoubleSpinBox()
        self.iou_threshold.setRange(0.01, 1.0)
        self.iou_threshold.setSingleStep(0.05)
        self.iou_threshold.setValue(0.45)
        params_layout.addRow("IoU阈值:", self.iou_threshold)
        
        # 类别过滤
        self.class_filter_edit = QLineEdit()
        self.class_filter_edit.setPlaceholderText("类别ID, 用逗号分隔, 留空为全部")
        params_layout.addRow("类别过滤:", self.class_filter_edit)
        
        left_layout.addWidget(params_group)
        
        # 操作按钮组
        actions_layout = QHBoxLayout()
        
        self.run_btn = QPushButton("运行推理")
        self.run_btn.clicked.connect(self.run_inference)
        
        self.stop_btn = QPushButton("停止")
        self.stop_btn.setEnabled(False)
        self.stop_btn.clicked.connect(self.stop_inference)
        
        self.save_btn = QPushButton("保存结果")
        self.save_btn.setEnabled(False)
        self.save_btn.clicked.connect(self.save_results)
        
        actions_layout.addWidget(self.run_btn)
        actions_layout.addWidget(self.stop_btn)
        actions_layout.addWidget(self.save_btn)
        
        left_layout.addLayout(actions_layout)
        
        # 进度条
        self.progress_bar = QProgressBar()
        self.progress_bar.setVisible(False)
        left_layout.addWidget(self.progress_bar)
        
        # 推理统计信息
        stats_group = QGroupBox("推理统计")
        stats_layout = QFormLayout(stats_group)
        
        self.inference_time_label = QLabel("0 ms")
        self.objects_detected_label = QLabel("0")
        self.image_size_label = QLabel("无")
        self.fps_label = QLabel("0")  # 添加FPS标签
        
        stats_layout.addRow("推理时间:", self.inference_time_label)
        stats_layout.addRow("检测到的目标:", self.objects_detected_label)
        stats_layout.addRow("图像尺寸:", self.image_size_label)
        stats_layout.addRow("FPS:", self.fps_label)  # 添加FPS显示
        
        left_layout.addWidget(stats_group)
        
        # 日志输出
        log_group = QGroupBox("日志")
        log_layout = QVBoxLayout(log_group)
        
        self.log_text = QTextEdit()
        self.log_text.setReadOnly(True)
        self.log_text.setMaximumHeight(100)
        log_layout.addWidget(self.log_text)
        
        left_layout.addWidget(log_group)
        
        # 添加弹性空间
        left_layout.addStretch()
        
        # 右侧结果显示面板
        right_panel = QWidget()
        right_layout = QVBoxLayout(right_panel)
        right_layout.setContentsMargins(0, 0, 0, 0)
        
        # 结果标签
        result_label = QLabel("推理结果")
        result_label.setAlignment(Qt.AlignCenter)
        result_label.setProperty("subtitle", "true")
        right_layout.addWidget(result_label)
        
        # 结果显示区域
        self.results_widget = ResultsWidget()
        right_layout.addWidget(self.results_widget)
        
        # 添加面板到分割器
        splitter.addWidget(left_panel)
        splitter.addWidget(right_panel)
        
        # 设置初始分割比例
        splitter.setSizes([int(self.width() * 0.3), int(self.width() * 0.7)])
        
        # 将分割器添加到主布局
        main_layout.addWidget(splitter)
        
    def browse_model(self):
        """浏览选择模型文件"""
        model_path, _ = QFileDialog.getOpenFileName(
            self, "选择YOLO模型文件", "", "PyTorch模型 (*.pt);;所有文件 (*.*)"
        )
        
        if model_path:
            self.model_path_edit.setText(model_path)
            
    def browse_source(self):
        """浏览选择输入源文件"""
        source_type = self.source_type_combo.currentText()
        
        if source_type == "图像":
            source_path, _ = QFileDialog.getOpenFileName(
                self, "选择输入图像", "", 
                "图像文件 (*.jpg *.jpeg *.png *.bmp);;所有文件 (*.*)"
            )
        elif source_type == "视频":
            source_path, _ = QFileDialog.getOpenFileName(
                self, "选择输入视频", "", 
                "视频文件 (*.mp4 *.avi *.mov *.mkv);;所有文件 (*.*)"
            )
        else:
            # 摄像头模式，直接使用摄像头ID
            self.source_path_edit.setText(self.camera_id_combo.currentText())
            return
        
        if source_path:
            self.source_path_edit.setText(source_path)
            
            # 如果是图像，预览图像
            if source_path.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp')):
                try:
                    img = cv2.imread(source_path)
                    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
                    
                    # 更新图像尺寸信息
                    h, w, c = img.shape
                    self.image_size_label.setText(f"{w}x{h}")
                    
                    # 显示原始图像
                    q_img = QImage(img.data, w, h, w * c, QImage.Format_RGB888)
                    pixmap = QPixmap.fromImage(q_img)
                    self.results_widget.setPixmap(pixmap.scaled(
                        self.results_widget.size(), 
                        Qt.KeepAspectRatio, 
                        Qt.SmoothTransformation
                    ))
                except Exception as e:
                    self.log_message(f"预览图像失败: {str(e)}", is_error=True)
            
            # 如果是视频，显示预览帧
            elif source_path.lower().endswith(('.mp4', '.avi', '.mov', '.mkv')):
                try:
                    cap = cv2.VideoCapture(source_path)
                    if cap.isOpened():
                        ret, frame = cap.read()
                        if ret:
                            # 获取视频信息
                            fps = cap.get(cv2.CAP_PROP_FPS)
                            width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
                            height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
                            frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
                            
                            self.image_size_label.setText(f"{width}x{height}")
                            self.fps_label.setText(f"{fps:.2f}")
                            
                            # 显示第一帧
                            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                            q_img = QImage(frame_rgb.data, width, height, width * 3, QImage.Format_RGB888)
                            pixmap = QPixmap.fromImage(q_img)
                            self.results_widget.setPixmap(pixmap.scaled(
                                self.results_widget.size(), 
                                Qt.KeepAspectRatio, 
                                Qt.SmoothTransformation
                            ))
                            
                            self.log_message(f"已加载视频: {os.path.basename(source_path)}, {width}x{height}, {fps}fps, {frames}帧")
                        cap.release()
                    else:
                        self.log_message(f"无法打开视频文件: {source_path}", is_error=True)
                except Exception as e:
                    self.log_message(f"预览视频失败: {str(e)}", is_error=True)
        
    def refresh_model_list(self):
        """刷新模型列表"""
        self.model_combo.clear()
        
        # 添加一个空选项
        self.model_combo.addItem("-- 选择模型 --")
        
        # 扫描模型目录
        scan_dirs = []
        
        # 1. 检查当前工作目录下的models目录
        models_dir = os.path.join(os.getcwd(), "models")
        if os.path.exists(models_dir):
            scan_dirs.append(models_dir)
            
        # 2. 检查项目目录下的models目录（如果有当前项目）
        from_main = self.parent()
        while from_main and not hasattr(from_main, 'current_project'):
            from_main = from_main.parent()
            
        if from_main and hasattr(from_main, 'current_project') and from_main.current_project:
            project_models_dir = os.path.join(from_main.current_project["path"], "models")
            if os.path.exists(project_models_dir):
                scan_dirs.append(project_models_dir)
                
        # 3. 检查.ultralytics/models目录（ultralytics默认下载目录）
        home_dir = os.path.expanduser("~")
        ultralytics_models_dir = os.path.join(home_dir, ".ultralytics", "models")
        if os.path.exists(ultralytics_models_dir):
            scan_dirs.append(ultralytics_models_dir)
            
        # 扫描所有目录中的.pt文件
        found_models = set()
        for directory in scan_dirs:
            for model_path in glob.glob(os.path.join(directory, "*.pt")):
                model_name = os.path.basename(model_path)
                if model_name not in found_models:
                    self.model_combo.addItem(model_name, model_path)
                    found_models.add(model_name)
        
    def on_model_selected(self, index):
        """当从下拉框选择模型时"""
        if index > 0:  # 跳过"-- 选择模型 --"选项
            model_path = self.model_combo.itemData(index)
            self.model_path_edit.setText(model_path)
            
    def on_source_type_changed(self, index):
        """当源类型发生变化时"""
        if index == 2:  # 摄像头
            self.source_path_edit.setReadOnly(True)
            self.source_path_edit.setText(self.camera_id_combo.currentText())
            self.camera_id_combo.setVisible(True)
            self.camera_label.setVisible(True)  # 使用引用的标签
            
            # 更新摄像头ID编辑框
            self.camera_id_combo.currentIndexChanged.connect(
                lambda idx: self.source_path_edit.setText(self.camera_id_combo.currentText())
            )
            
            # 重命名按钮
            self.browse_source_btn.setText("测试连接")
            try:
                self.browse_source_btn.clicked.disconnect()
            except:
                pass
            self.browse_source_btn.clicked.connect(self.test_camera_connection)
        else:
            self.source_path_edit.setReadOnly(True)
            self.source_path_edit.setText("")
            self.camera_id_combo.setVisible(False)
            self.camera_label.setVisible(False)  # 使用引用的标签
            
            # 恢复按钮功能
            self.browse_source_btn.setText("浏览...")
            try:
                self.browse_source_btn.clicked.disconnect()
            except:
                pass
            self.browse_source_btn.clicked.connect(self.browse_source)
    
    def test_camera_connection(self):
        """测试摄像头连接"""
        camera_id = int(self.camera_id_combo.currentText())
        
        try:
            cap = cv2.VideoCapture(camera_id)
            if not cap.isOpened():
                self.log_message(f"无法连接到摄像头 ID: {camera_id}", is_error=True)
                return
                
            # 读取一帧
            ret, frame = cap.read()
            if not ret:
                self.log_message(f"无法从摄像头 ID: {camera_id} 读取图像", is_error=True)
                cap.release()
                return
                
            # 获取摄像头信息
            width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
            height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
            
            # 显示图像
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            q_img = QImage(frame_rgb.data, width, height, width * 3, QImage.Format_RGB888)
            pixmap = QPixmap.fromImage(q_img)
            self.results_widget.setPixmap(pixmap.scaled(
                self.results_widget.size(), 
                Qt.KeepAspectRatio, 
                Qt.SmoothTransformation
            ))
            
            # 更新信息
            self.image_size_label.setText(f"{width}x{height}")
            self.log_message(f"成功连接到摄像头 ID: {camera_id}, 分辨率: {width}x{height}")
            
            # 释放摄像头
            cap.release()
            
        except Exception as e:
            self.log_message(f"测试摄像头连接失败: {str(e)}", is_error=True)
    
    def run_inference(self):
        """运行推理"""
        # 检查模型路径
        model_path = self.model_path_edit.text()
        if not model_path:
            QMessageBox.warning(self, "警告", "请先选择模型文件")
            return
            
        # 检查源文件路径
        source_path = self.source_path_edit.text()
        if not source_path:
            QMessageBox.warning(self, "警告", "请先选择输入源文件或摄像头")
            return
            
        # 获取源类型
        source_type_text = self.source_type_combo.currentText()
        if source_type_text == "图像":
            source_type = "image"
        elif source_type_text == "视频":
            source_type = "video"
        else:
            source_type = "camera"
            
        # 获取参数
        conf_thres = self.conf_threshold.value()
        iou_thres = self.iou_threshold.value()
        
        # 解析类别过滤
        classes = None
        if self.class_filter_edit.text().strip():
            try:
                classes = [int(c.strip()) for c in self.class_filter_edit.text().split(',')]
            except ValueError:
                QMessageBox.warning(self, "警告", "类别过滤格式错误，请使用逗号分隔的数字")
                return
                
        # 更新UI状态
        self.run_btn.setEnabled(False)
        self.stop_btn.setEnabled(True)
        self.progress_bar.setValue(0)
        self.progress_bar.setVisible(True)
        
        # 清空统计信息
        self.inference_time_label.setText("计算中...")
        self.objects_detected_label.setText("计算中...")
        self.fps_label.setText("计算中...")
        
        # 创建并启动推理线程
        self.inference_thread = InferenceThread(
            model_path, source_path, conf_thres, iou_thres, classes, source_type
        )
        self.inference_thread.progress_updated.connect(self.update_progress)
        self.inference_thread.result_ready.connect(self.display_results)
        self.inference_thread.video_frame.connect(self.display_video_frame)
        self.inference_thread.inference_finished.connect(self.inference_completed)
        self.inference_thread.inference_stats.connect(self.update_stats)
        
        self.log_message(f"开始推理: 模型={os.path.basename(model_path)}, 源={os.path.basename(source_path)}, 类型={source_type_text}")
        self.inference_thread.start()
        
    def stop_inference(self):
        """停止推理"""
        if self.inference_thread and self.inference_thread.isRunning():
            self.inference_thread.stop()
            self.inference_thread.wait()
            self.log_message("推理已停止")
            
        self.run_btn.setEnabled(True)
        self.stop_btn.setEnabled(False)
        
    def update_progress(self, value):
        """更新进度条"""
        self.progress_bar.setValue(value)
        
    def display_results(self, results, orig_image):
        """显示推理结果(用于图像)"""
        self.results_widget.set_results(results, orig_image)
        self.save_btn.setEnabled(True)
        
    def display_video_frame(self, results, frame):
        """显示视频帧结果"""
        self.results_widget.set_results(results, frame)
        # 视频模式下也允许保存当前帧
        self.save_btn.setEnabled(True)
        
    def inference_completed(self, success, message):
        """推理完成的回调"""
        self.run_btn.setEnabled(True)
        self.stop_btn.setEnabled(False)
        
        if success:
            self.log_message(message)
        else:
            self.log_message(message, is_error=True)
            
        # 隐藏进度条
        self.progress_bar.setVisible(False)
        
    def update_stats(self, stats):
        """更新统计信息"""
        self.inference_time_label.setText(f"{stats.get('inference_time', 0):.3f} 秒")
        self.objects_detected_label.setText(str(stats.get('detected_objects', 0)))
        self.image_size_label.setText(stats.get('image_size', '无'))
        
        # 显示FPS（如果有）
        if 'fps' in stats:
            self.fps_label.setText(f"{stats['fps']:.2f}")
        else:
            self.fps_label.setText("N/A")
        
    def save_results(self):
        """保存推理结果图像"""
        if not hasattr(self.results_widget, 'results') or self.results_widget.results is None:
            QMessageBox.warning(self, "警告", "没有可保存的结果")
            return
            
        # 获取当前源类型
        source_type = self.source_type_combo.currentText()
        
        # 获取保存路径
        source_path = self.source_path_edit.text()
        
        # 构建默认文件名
        if source_type == "摄像头":
            default_filename = f"camera_{int(time.time())}_result.jpg"
            default_path = os.path.join(os.getcwd(), default_filename)
        else:
            source_name = os.path.basename(source_path)
            base_name, ext = os.path.splitext(source_name)
            default_path = os.path.join(os.path.dirname(source_path), f"{base_name}_result.jpg")
        
        save_path, _ = QFileDialog.getSaveFileName(
            self, "保存结果图像", 
            default_path,
            "图像文件 (*.jpg *.jpeg *.png)"
        )
        
        if save_path:
            try:
                # 使用ultralytics的保存方法
                result_img = self.results_widget.results.plot()
                cv2.imwrite(save_path, result_img)
                self.log_message(f"结果已保存到: {save_path}")
            except Exception as e:
                self.log_message(f"保存结果失败: {str(e)}", is_error=True)
                
    def log_message(self, message, is_error=False):
        """添加日志消息"""
        # 设置文本颜色
        color = "red" if is_error else "black"
        self.log_text.append(f'<span style="color: {color};">{message}</span>')
        
        # 滚动到底部
        self.log_text.verticalScrollBar().setValue(
            self.log_text.verticalScrollBar().maximum()
        )
        
    def closeEvent(self, event):
        """窗口关闭事件"""
        # 确保线程正确终止
        if self.inference_thread and self.inference_thread.isRunning():
            self.inference_thread.stop()
            self.inference_thread.wait()
            
        super().closeEvent(event) 